#include <Windows.h>
#include <Vfw.h>
#include "Core.h"

map<HWND, Core*> Core::instances;

Core::Core(HINSTANCE hInstance)
{
	active = false;
	locked = false;
	lockAtUnplug = false;	

	activeChangedPtr = 0;
	lockedChangedPtr = 0;

	hiddenWindow = new Window();
	hiddenWindow->Create(hInstance, TEXT("AuroraWindow"), HiddenWindowWndProc, 0, WS_POPUP, CS_NOCLOSE);
	instances[hiddenWindow->GetHWnd()] = this;

	lockWindow = new Window();
	lockWindow->Create(hInstance,
					  TEXT("AuroraWindow"),
					  LockWindowWndProc,
					  WS_EX_TOOLWINDOW | WS_EX_TOPMOST,
					  WS_POPUP,
					  CS_NOCLOSE,
					  GetStockBrush(BLACK_BRUSH)/*doesn't work?!*/,
					  0,
					  LoadIcon(hInstance, MAKEINTRESOURCE(IDI_ICON1)));

	//set bg brush again due to the problem above
	lockWindow->SetBackgroundBrush(GetStockBrush(BLACK_BRUSH));

	lockWindow->SetLocation(GetSystemMetrics(SM_XVIRTUALSCREEN), GetSystemMetrics(SM_YVIRTUALSCREEN));
	lockWindow->SetSize(GetSystemMetrics(SM_CXVIRTUALSCREEN), GetSystemMetrics(SM_CYVIRTUALSCREEN));


	WTSRegisterSessionNotification(hiddenWindow->GetHWnd(), NOTIFY_FOR_THIS_SESSION);
	RegisterPowerSettingNotification(hiddenWindow->GetHWnd(), &GUID_LIDSWITCH_STATE_CHANGE, DEVICE_NOTIFY_WINDOW_HANDLE);
	RefreshKeyDevicesVolumeLetters();
}

Core::~Core()
{
	WTSUnRegisterSessionNotification(hiddenWindow->GetHWnd());

	instances.erase(hiddenWindow->GetHWnd());

	delete lockWindow;
	lockWindow = 0;

	delete hiddenWindow;
	hiddenWindow = 0;
}


bool Core::Activate(bool lock)
{
	if (active)
		return false;

	if (keyDevices.empty())
		return false;

	active = true;
	OnActiveChanged(active);

	if (lock || (lockAtUnplug && !IsAnyKeyDevicePluggedIn()))
		Lock();

	return true;
}

void Core::Deactivate()
{
	if (!active)
		return;

	active = false;
	OnActiveChanged(active);
}


bool Core::IsActive()
{
	return active;
}

bool Core::IsLocked()
{
	return locked;
}

void Core::SetLockAtUnplug(bool value)
{
	lockAtUnplug = value;	
	OnLockAtUnplugChanged(lockAtUnplug);

	if (active && lockAtUnplug)
	{
		if (!IsAnyKeyDevicePluggedIn())
			Lock();
	}
}

bool Core::GetLockAtUnplug()
{
	return lockAtUnplug;
}


void Core::SetKeyDeviceCollection(map<DWORD, tstring> collection)
{
	keyDevices = collection;
	RefreshKeyDevicesVolumeLetters();
	OnKeyDevicesChanged(GetKeyDeviceCount());
}

map<DWORD, tstring> Core::GetKeyDeviceCollection()
{
	return keyDevices;
}


int Core::RemoveKeyDevice(DWORD serial)
{
	if (!serial)
		return KEYDEVICE_INVALIDSERIAL;
	else if (keyDevices.find(serial) == keyDevices.end())
		return KEYDEVICE_NOTFOUND;

	keyDevices.erase(serial);
	RefreshKeyDevicesVolumeLetters();
	
	if (GetKeyDeviceCount() > 0)
	{
		if (active && lockAtUnplug && !IsAnyKeyDevicePluggedIn())
			Lock();
	}
	else
	{
		Deactivate();
	}

	OnKeyDevicesChanged(GetKeyDeviceCount());
	
	return KEYDEVICE_SUCCESS;
}

int Core::AddKeyDevice(DWORD serial, tstring name)
{
	if (!serial)
		return KEYDEVICE_INVALIDSERIAL;
	else if (keyDevices.find(serial) != keyDevices.end())
		return KEYDEVICE_ALREADYEXISTS;

	keyDevices[serial] = name;
	RefreshKeyDevicesVolumeLetters();
	OnKeyDevicesChanged(GetKeyDeviceCount());

	return KEYDEVICE_SUCCESS;
}

void Core::ClearKeyDevices()
{
	keyDevices.clear();
	RefreshKeyDevicesVolumeLetters();
	OnKeyDevicesChanged(GetKeyDeviceCount());
	Deactivate();
}

size_t Core::GetKeyDeviceCount()
{
	return keyDevices.size();
}


void Core::ActiveChanged(STATECHANGEDPROC ptr)
{
	this->activeChangedPtr = ptr;
}

void Core::LockedChanged(STATECHANGEDPROC ptr)
{
	this->lockedChangedPtr = ptr;
}

void Core::LockAtUnplugChanged(STATECHANGEDPROC ptr)
{
	this->lockAtUnplugChangedPtr = ptr;
}

void Core::KeyDevicesChanged(SIZECHANGEDPROC ptr)
{
	this->keyDevicesChangedPtr = ptr;
}

void Core::UnauthorizedAccess(EVENTPROC ptr)
{
	this->unauthorizedAccessPtr = ptr;
}


void Core::Lock()
{
	if (!active)
		return;

	if (keyDevices.empty())
		return;

	if (!locked)
	{
		locked = true;
		OnLockedChanged(locked);
	}

	LockWorkStation();
}

void Core::LockActions()
{
	PostMessage(hiddenWindow->GetHWnd(), WM_SYSCOMMAND, SC_SCREENSAVE, 0);
}

void Core::UnauthorizedAccess()
{
	lockWindow->SetLocation(GetSystemMetrics(SM_XVIRTUALSCREEN), GetSystemMetrics(SM_YVIRTUALSCREEN));
	lockWindow->SetSize(GetSystemMetrics(SM_CXVIRTUALSCREEN), GetSystemMetrics(SM_CYVIRTUALSCREEN));
	lockWindow->Show();
	lockWindow->SetTopMost(true);

	OnUnauthorizedAccess();

	//lock again:
	Lock();
}

void Core::Unlock()
{
	lockWindow->Hide();	

	if (locked)
	{
		locked = false;
		OnLockedChanged(locked);
	}
}


void Core::RefreshKeyDevicesVolumeLetters()
{
    keyDeviceVolumeLetters.clear();

	if (keyDevices.empty())
        return;

	set<TCHAR> drives = Utilities::GetDrives();

	set<TCHAR>::iterator iter;
	for (iter = drives.begin(); iter != drives.end(); iter++)
	{
        TCHAR driveRoot[] = {*iter, ':', '\\', '\0'};
		int driveType = GetDriveType(driveRoot);

		if (driveType == DRIVE_FIXED || driveType == DRIVE_REMOVABLE)
		{
			DWORD serial = Utilities::GetVolumeSerial(*iter);

			if (serial && keyDevices.find(serial) != keyDevices.end())
			{
				keyDeviceVolumeLetters.insert(*iter);
			}
		}
	}
}

bool Core::IsAnyKeyDevicePluggedIn()
{
	set<TCHAR> drives = Utilities::GetDrives();

	set<TCHAR>::iterator iter;
	for (iter = drives.begin(); iter != drives.end(); iter++)
	{
		TCHAR driveRoot[] = {*iter, ':', '\\', '\0'};
		int driveType = GetDriveType(driveRoot);

		if (driveType == DRIVE_FIXED || driveType == DRIVE_REMOVABLE)
		{
			DWORD serial = Utilities::GetVolumeSerial(*iter);

			if (serial && keyDevices.find(serial) != keyDevices.end())
				return true;
		}

	}


    return false;
}


void Core::OnActiveChanged(bool value)
{
	if (activeChangedPtr)
		activeChangedPtr(this, value);
}

void Core::OnLockedChanged(bool value)
{
	if (lockedChangedPtr)
		lockedChangedPtr(this, value);
}

void Core::OnLockAtUnplugChanged(bool value)
{
	if (lockAtUnplugChangedPtr)
		lockAtUnplugChangedPtr(this, value);
}

void Core::OnKeyDevicesChanged(size_t value)
{
	if (keyDevicesChangedPtr)
		keyDevicesChangedPtr(this, value);
}

void Core::OnUnauthorizedAccess()
{
	if (unauthorizedAccessPtr)
		unauthorizedAccessPtr(this);
}


LRESULT CALLBACK Core::HiddenWindowWndProc(LPVOID sender, HWND hWnd, UINT uMsg, WPARAM wParam, LPARAM lParam)
{	
	Core* instance = instances[hWnd];

	switch (uMsg)
	{
	case WM_WTSSESSION_CHANGE:
		{
			if (instance->active)
			{
				if (wParam == WTS_SESSION_UNLOCK)
				{
					if (instance->IsAnyKeyDevicePluggedIn())
						instance->Unlock();
					else
						instance->UnauthorizedAccess();
				}
				else if (wParam == WTS_SESSION_LOCK)
				{
					instance->LockActions();
				}
			}

			break;
		}
	case WM_DEVICECHANGE:
		{
			if (wParam == DBT_DEVICEARRIVAL || wParam == DBT_DEVICEREMOVECOMPLETE)
			{
				DEV_BROADCAST_HDR *hdr = (DEV_BROADCAST_HDR*)lParam;

				if (hdr->dbch_devicetype == DBT_DEVTYP_VOLUME)
				{
					DEV_BROADCAST_VOLUME *volume = (DEV_BROADCAST_VOLUME*)lParam;

					if (wParam == DBT_DEVICEARRIVAL)
					{	
						TCHAR volumeLetter = Utilities::GetVolumeLetter(volume->dbcv_unitmask);
						DWORD serial = Utilities::GetVolumeSerial(volumeLetter);

						instance->RefreshKeyDevicesVolumeLetters();
					}
					else if (wParam == DBT_DEVICEREMOVECOMPLETE)
					{
						TCHAR volumeLetter = Utilities::GetVolumeLetter(volume->dbcv_unitmask);

						if (instance->active && instance->lockAtUnplug && instance->keyDeviceVolumeLetters.find(volumeLetter) != instance->keyDeviceVolumeLetters.end() && !instance->IsAnyKeyDevicePluggedIn())
						{
							//unplugged device is a key and no other key is still plugged-in
							instance->Lock();
						}
					}
				}
			}

			break;
		}
	case WM_POWERBROADCAST:
		{
			if (wParam == PBT_POWERSETTINGCHANGE && instance->active && instance->locked)
			{
				POWERBROADCAST_SETTING* pbs = (POWERBROADCAST_SETTING*)lParam;

				if (pbs->PowerSetting == GUID_LIDSWITCH_STATE_CHANGE && pbs->DataLength == sizeof(int))
				{
					int status = *(int*)(DWORD_PTR) pbs->Data;

					if (status == 0)
					{
						//lid is closed
					}
					else if (status == 1)
					{
						//lid is opened
					}
				}
			}
			break;
		}
	}

	return DefWindowProc(hWnd, uMsg, wParam, lParam);
}

LRESULT CALLBACK Core::LockWindowWndProc(LPVOID sender, HWND hWnd, UINT uMsg, WPARAM wParam, LPARAM lParam)
{	
	Core* instance = instances[hWnd];

	switch (uMsg)
	{
	case WM_PAINT:
		{
			RECT clientRect;
			GetClientRect(hWnd, &clientRect);

			HDC          hdc;
			PAINTSTRUCT  ps;

			hdc = BeginPaint(hWnd, &ps);

			Utilities::DrawDefaultText(clientRect, hdc, TEXT("Unauthorized Access"), 36, RGB(0xbb, 0x04, 0x0b), DT_LEFT | DT_TOP | DT_SINGLELINE, FW_BOLD, 80, 80);
			Utilities::DrawDefaultText(clientRect, hdc, TEXT("Unauthorized Access"), 36, RGB(0xbb, 0x04, 0x0b), DT_RIGHT | DT_BOTTOM | DT_SINGLELINE, FW_BOLD, -80, -80);

			EndPaint(hWnd, &ps);

			break;
		}
	}

	return DefWindowProc(hWnd, uMsg, wParam, lParam);
}



//-----experimental-----

//photo (*.dib)
void Core::CaptureImage(HWND parentHWnd)
{
	//check for right webcam index:

	//TCHAR buffer[256] = {0};
	//TCHAR version[256] = {0};
	//capGetDriverDescription(0, 0, 256, 0, 256);
	//onnly valid if a version and a description in buffer is returned

	//HWND h2 = capCreateCaptureWindow(TEXT("camCaptureWnd"), WS_VISIBLE | WS_CHILD, 0, 0, 20, 20, parentHWnd, 0);

	//if (capDriverConnect(h2, 0))
	//{
	//	time_t now = time(0);

	//	tm timeInfo;
	//	localtime_s(&timeInfo, &now);

	//	//solve the DST problem
	//	if (timeInfo.tm_isdst > 0)
	//	{
	//		timeInfo.tm_hour++;
	//		now = mktime(&timeInfo);
	//	}

	//	TCHAR fileName[_MAX_FNAME] = {0};
	//	_stprintf_s(fileName, TEXT("Capture_%d-%02d-%02d_%02d-%02d-%02d.avi"), timeInfo.tm_year, timeInfo.tm_mon, timeInfo.tm_wday, timeInfo.tm_hour, timeInfo.tm_min, timeInfo.tm_sec);
	//	
	//	capFileSaveDIB(h2, fileName);
	//	capCaptureStop(h2);
	//}
}

//7s video (*.avi)
void Core::CaptureVideo(HWND parentHWnd)
{
	//check for right webcam index:

	//TCHAR buffer[256] = {0};
	//TCHAR version[256] = {0};
	//capGetDriverDescription(0, 0, 256, 0, 256);
	//onnly valid if a version and a description in buffer is returned

	//HWND h2 = capCreateCaptureWindow(TEXT("camCaptureWnd"), WS_VISIBLE | WS_CHILD, 0, 0, 20, 20, parentHWnd, 0);

	//if (capDriverConnect(h2, 0))
	//{
	//	CAPTUREPARMS cp = {0};
	//	capCaptureGetSetup(h2, &cp, sizeof(cp));
	//	cp.fYield = TRUE;
	//	cp.fLimitEnabled = TRUE;
	//	cp.wTimeLimit = 7;
	//	cp.vKeyAbort = 0;
	//	cp.fAbortLeftMouse = FALSE;
	//	cp.fAbortRightMouse = FALSE;
	//	capCaptureSetSetup(h2, &cp, sizeof(cp));


	//	time_t now = time(0);

	//	tm timeInfo;
	//	localtime_s(&timeInfo, &now);

	//	//solve the DST problem
	//	if (timeInfo.tm_isdst > 0)
	//	{
	//		timeInfo.tm_hour++;
	//		now = mktime(&timeInfo);
	//	}

	//	TCHAR fileName[_MAX_FNAME] = {0};
	//	_stprintf_s(fileName, TEXT("Capture_%d-%02d-%02d_%02d-%02d-%02d.avi"), timeInfo.tm_year, timeInfo.tm_mon, timeInfo.tm_wday, timeInfo.tm_hour, timeInfo.tm_min, timeInfo.tm_sec);
	//	capFileSetCaptureFile(h2, fileName);

	//	capCaptureSequence(h2);
	//	
	//	//call this later!
	//	capCaptureStop(h2);
	//}
}

